# coding=utf8
import numpy as np
import libs.NavieBayes as naiveBayes
from libs.help import get_root_path


def main():
    base_path = get_root_path('data')
    # 获取语料
    filename = base_path + 'training/AdCollection.txt'
    ad_words, class_lables = naiveBayes.loadSMSData(filename)
    vocabularyList = naiveBayes.createVocabularyList(ad_words)
    print("生成语料库！")
    trainMarkedWords = naiveBayes.setOfWordsListToVecTor(vocabularyList, ad_words)
    print("数据标记完成！")
    # 转成array向量
    trainMarkedWords = np.array(trainMarkedWords)
    print("数据转成矩阵！")
    pWordsSpamicity, pWordsHealthy, pSpam = naiveBayes.trainingNaiveBayes(trainMarkedWords, class_lables)
    print('pSpam:', pSpam)
    with open(base_path + 'trained/pSpam.txt', 'w', encoding='utf8') as f:
        f.write(str(pSpam))
    # 保存训练生成的语料库信息
    # 保存语料库词汇
    with open(base_path + 'trained/vocabularyList.txt', 'w', encoding='utf8')as f:
        f.write('\t'.join(vocabularyList))
    # 保存训练阶段获取的参数：pWordsSpamicity和pWordsHealthy
    np.savetxt(base_path + 'trained/pWordsSpamicity.txt', pWordsSpamicity, delimiter='\t')
    np.savetxt(base_path + 'trained/pWordsHealthy.txt', pWordsHealthy, delimiter='\t')


if __name__ == '__main__':
    main()
